{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  import pydrake\n",
    "  import underactuated\n",
    "except ImportError:\n",
    "  !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "  from jupyter_setup import setup_underactuated\n",
    "  setup_underactuated()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Autapse\n",
    "\n",
    "The simplest recurrent neural network model.  $$\\dot{x} = -x + \\tanh(wx + u)$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from ipywidgets import interact\n",
    "\n",
    "from underactuated.jupyter import SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "\n",
    "def autapse(x, w=1, u=0):\n",
    "    \"\"\"Args:\n",
    "    w is feedback weight\n",
    "    u is input\n",
    "    \"\"\"\n",
    "    return -x + np.tanh(w * x + u)\n",
    "\n",
    "\n",
    "Autapse = np.vectorize(autapse)\n",
    "xmax = 2.\n",
    "ymax = 1.\n",
    "x = np.arange(-xmax, xmax, 0.01)\n",
    "\n",
    "fig = plt.figure()\n",
    "\n",
    "if plt_is_interactive:\n",
    "  line, = plt.plot(x, Autapse(x, w=3., u=0), linewidth=2.)\n",
    "\n",
    "  def update(w=1, u=0):\n",
    "    line.set_ydata(Autapse(x, w=w, u=u))\n",
    "    fig.canvas.draw()\n",
    "    \n",
    "  interact(update, w=(0,3,0.1), u=(-1.5,1.5,0.1))\n",
    "  \n",
    "else:\n",
    "  plt.plot(x, Autapse(x, w=.75, u=0), linewidth=2., label=\"w=3/4\")\n",
    "  plt.plot(x, Autapse(x, w=3., u=0), linewidth=2., label=\"w=3\")\n",
    "  plt.legend()\n",
    "\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylim((-ymax, ymax))\n",
    "plt.ylabel(\"xdot\")\n",
    "\n",
    "# draw the x and y axes.\n",
    "plt.plot([-xmax, xmax], [0, 0], color=\"k\", linestyle=\"-\", linewidth=1.)\n",
    "plt.plot([0, 0], [-ymax, ymax], color=\"k\", linestyle=\"-\", linewidth=1.)\n",
    "# draw the line through the origin with slope -1.\n",
    "plt.plot([-ymax, ymax], [ymax, -ymax], color=\"k\", linestyle=\"-\", linewidth=1.)\n",
    "\n",
    "plt.title(\"Autapse\")\n",
    "plt.axis(\"equal\");\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Long short-term memory (LSTM)\n",
    "\n",
    "A recurrent neural network component..  [Wikipedia](https://en.wikipedia.org/wiki/Long_short-term_memory)\n",
    "\n",
    "We'll look at the [\"gated recurrent unit\" version](https://colah.github.io/posts/2015-08-Understanding-LSTMs/)...\n",
    "\\begin{align}\n",
    "i[n] &= \\sigma_g(A_{i} x[n] + B_{i} u[n] + c_i) & \\text{input gate} \\\\\n",
    "o[n] &= \\sigma_g(A_{o} x[n] + B_{o} u[n] + c_o) & \\text{output gate} \\\\\n",
    "f[n] &= \\sigma_g(A_{f} x[n] + B_{f} u[n] + c_f) & \\text{forget gate} \\\\\n",
    "x[n+1] &= f[n] \\circ x[n] + i[n] \\circ \\tanh(B_{x} u[n] + c_c) \\\\\n",
    "y[n] &= \\tanh(o[n] \\circ x[n])\n",
    "\\end{align}\n",
    "\n",
    "where the operator $\\circ$  denotes the Hadamard product (element-wise product). $\\sigma _{g}(x) = \\frac{1}{1+e^{-x}} \\in (0,1)$ is the sigmoid function and recall that $\\tanh(x) \\in (-1,1)$.\n",
    "\n",
    "In the example below, we will plot the continuous-time version of this (to stay consistent with the rest of the analysis in the chapter), $\\dot{x}$ vs $x$.  To keep the number of parameters reasonable, I've assumed that there are three inputs -- with $$u = \\begin{bmatrix} u_{input} \\\\ u_{forget} \\\\ u_{x} \\end{bmatrix}, B_i = [ 1, 0 0 ], B_f = [ 0, 1, 0 ], B_x = [ 0, 0, 1].$$ Note that the output $y$ has no impact on the stability and is not represented below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipywidgets import interact\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "def sigma(x):\n",
    "  return 1./(1+np.exp(-x))\n",
    "\n",
    "def lstm(x, uf=0, ui=0, ux=0, af=1, bf=1, cf=0, ai=1, bi=1, ci=0, bx=1, cx=0):\n",
    "  return - x + sigma(af*x + bf*uf + cf)*x + sigma(ai*x + bi*ui + ci) * np.tanh(bx*ux+cx)\n",
    "\n",
    "Lstm = np.vectorize(lstm)\n",
    "xmax = 10.\n",
    "ymax = 4.\n",
    "x = np.arange(-xmax, xmax, 0.01)\n",
    "\n",
    "fig = plt.figure()\n",
    "\n",
    "if plt_is_interactive:\n",
    "  line, = plt.plot(x, Lstm(x, uf=-10, ui=-10, ux=0), linewidth=2.)\n",
    "\n",
    "  def update(u_forget, u_input, u_x,\n",
    "             a_forget, b_forget, c_forget,\n",
    "             a_input, b_input, c_input,\n",
    "             b_x, c_x):\n",
    "    line.set_ydata(Lstm(x,\n",
    "                        uf=u_forget, ui=u_input, ux=u_x,\n",
    "                        af=a_forget, bf=b_forget, cf=c_forget,\n",
    "                        ai=a_input, bi=b_input, ci=c_input,\n",
    "                        bx=b_x, cx=c_x))\n",
    "    fig.canvas.draw()\n",
    "    \n",
    "  interact(update,\n",
    "           u_forget=(-10,10,0.1),\n",
    "           u_input=(-10,10,0.1),\n",
    "           u_x=(-10,10,0.1),\n",
    "           a_forget=(0,2,0.1),\n",
    "           b_forget=(0,2,0.1),\n",
    "           c_forget=(-1,1,.1),\n",
    "           a_input=(0,2,0.1),\n",
    "           b_input=(0,2,0.1),\n",
    "           c_input=(-1,1,.1),\n",
    "           b_x=(0,2,0.1),\n",
    "           c_x=(-1,1,.1),\n",
    "          )\n",
    "  \n",
    "else:\n",
    "  def plot_lstm(uf=0, ui=0, ux=0):\n",
    "    plt.plot(x, Lstm(x, uf=uf, ui=ui, ux=ux), linewidth=2., label=\"uf=\" + str(uf) + \",ui=\" + str(ui) + \", ux=\" + str(ux))\n",
    "  \n",
    "  for us in np.arange(-1,1,.25):\n",
    "    plot_lstm(uf=0, ui=0, ux=us)\n",
    "  plt.legend()\n",
    "\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylim((-ymax, ymax))\n",
    "plt.ylabel(\"xdot\")\n",
    "\n",
    "# draw the x and y axes.\n",
    "plt.plot([-xmax, xmax], [0, 0], color=\"k\", linestyle=\"-\", linewidth=1.)\n",
    "plt.plot([0, 0], [-ymax, ymax], color=\"k\", linestyle=\"-\", linewidth=1.)\n",
    "# draw the line through the origin with slope -1.\n",
    "plt.plot([-ymax, ymax], [ymax, -ymax], color=\"k\", linestyle=\"-\", linewidth=1.)\n",
    "\n",
    "plt.title(\"LSTM\")\n",
    "plt.axis(\"equal\");\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I will make a few observations.  First, I'll admit that the behavior here looks pretty strange to me!  When \"forget gate\" is on (corresponding to a large negative input $u_{forget}$) and the \"input gate\" is off (corresponding to a large negative input $u_{input}$), the system has a stable fixed point close to zero.  When the forget gate is off (large $u_{forget}$ -- yes, it feels backwards to me, too) then the system becomes an integrator which will accumulate the input $u_x$.  Yes -- it's unstable!  And no, it's not an artifact of the continuous-time implementation.  The total output of the system is limited by the output $\\tanh$, but the internal state is unstable.\n",
    "\n",
    "The design of these units does not appear to have been made by someone who knew dynamics (nor graphical analysis). That they seem to work somewhat well in practice is fairly remarkable to me.  Shouldn't we be able to do better?\n",
    "\n",
    "Note: The parameter space here is large and the interpretations subtle.  If you find a better parameter regime and/or different interpretation, then I'd love to hear about it."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
